package org.col.dspace.tagger;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.util.Collections;

import org.nzdl.kea.main.KEAModelBuilder;
import org.nzdl.kea.stemmers.PorterStemmer;
import org.nzdl.kea.stopwords.StopwordsEnglish;


public class ModelBuilder {
	
	private KEAModelBuilder km;
	private void setOptionsTraining() {
		
		km = new KEAModelBuilder();
		
		// A. required arguments (no defaults):
		
		// 1. Name of the directory -- give the path to your directory with documents and keyphrases
		//    documents should be in txt format with an extention "txt"
		//    keyphrases with the same name as documents, but extension "key"
		//    one keyphrase per line!
		
		km.setDirName("en/traindocs");
		
		// 2. Name of the model -- give the path to where the model is to be stored and its name
		//String model_file=Thread.currentThread().getContextClassLoader().getResource("").getFile().concat("en/model");
		String model_file=Thread.currentThread().getContextClassLoader().getResource("").getFile().concat("en/model");
		km.setModelName(model_file);
		 
		// 3. Name of the vocabulary -- name of the file (without extension) that is stored in VOCABULARIES
		//    or "none" if no Vocabulary is used (free keyphrase extraction).
		km.setVocabulary("agrovoc");
		
		// 4. Format of the vocabulary in 3. Leave empty if vocabulary = "none", use "skos" or "txt" otherwise.
		km.setVocabularyFormat("skos");
		
//		 B. optional arguments if you want to change the defaults
		// 5. Encoding of the document
		km.setEncoding("UTF-8");
		
		// 6. Language of the document -- use "es" for Spanish, "fr" for French
		//    or other languages as specified in your "skos" vocabulary 
		km.setDocumentLanguage("en"); // es for Spanish, fr for French
		
		// 7. Stemmer -- adjust if you use a different language than English or if you want to alterate results
		// (We have obtained better results for Spanish and French with NoStemmer) 		
		km.setStemmer(new PorterStemmer()); 
		
		// 8. Stopwords -- adjust if you use a different language than English!
		km.setStopwords(new StopwordsEnglish());
		
		// 9. Maximum length of a keyphrase
		km.setMaxPhraseLength(3);
		
		// 10. Minimum length of a keyphrase
		km.setMinPhraseLength(1);
		
		// 11. Minumum occurrence of a phrase in the document -- use 2 for long documents!
		km.setMinNumOccur(2);
	
	//  Optional: turn off the keyphrase frequency feature	
	//	km.setUseKFrequency(false);
		
	}

	private void createModel() {
		try {
			km.buildModel(km.collectStems());

			km.saveModel();
			} catch (Exception e) {
				e.printStackTrace();
			}
	
}
	
	public static void main(String[] params) throws IOException{
		
		System.out.println("creating model from traning docs");

		ModelBuilder model = new ModelBuilder();
		
		model.setOptionsTraining();
		model.createModel();
		
		System.out.println("finished creating the model");

		
	}

}
